# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Train a YOLOv5 model on a custom dataset.

Models and datasets download automatically from the latest YOLOv5 release.
Models: https://github.com/ultralytics/yolov5/tree/master/models
Datasets: https://github.com/ultralytics/yolov5/tree/master/data
Tutorial: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

Usage:
    $ python path/to/train.py --data coco128.yaml --weights yolov5s.pt --img 640  # from pretrained (RECOMMENDED)
    $ python path/to/train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640  # from scratch
"""
import argparse
import math
import os
import pickle
import random
import sys
import time
from copy import deepcopy
from datetime import datetime
from pathlib import Path

import numpy as np
import torch
import torch.distributed as dist
import yaml
from torch.optim import lr_scheduler
from tqdm import tqdm

import sys

sys.path.append("C:/Users/cyc/Desktop/ppq/ppq/ppq/samples/Imagenet")

from trainer import ImageNetTrainer
from ppq.api import ENABLE_CUDA_KERNEL, export_ppq_graph
import ppq.lib as PFL
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
import torch
import torchvision
from ppq.api import ENABLE_CUDA_KERNEL, load_native_graph, load_torch_model, export, load_onnx_graph
from ppq.core import TargetPlatform
from ppq.executor import TorchExecutor
from ppq.quantization.optim import *
from ppq.quantization.quantizer import TensorRTQuantizer, FPGAQuantizer

# from Utilities.Imagenet import *  # check ppq.samples.Imagenet.Utilities
# from Utilities.Imagenet.imagenet_util import \
#     load_imagenet_from_directory  # check ppq.samples.Imagenet.Utilities

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

import val  # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yolo import Model
from utils.autoanchor import check_anchors
from utils.autobatch import check_train_batch_size
from utils.callbacks import Callbacks
from utils.dataloaders import create_dataloader
from utils.downloads import attempt_download, is_url
from utils.general import (LOGGER, check_amp, check_dataset, check_file, check_git_status, check_img_size,
                           check_requirements, check_suffix, check_yaml, colorstr, get_latest_run, increment_path,
                           init_seeds, intersect_dicts, labels_to_class_weights, labels_to_image_weights, methods,
                           one_cycle, print_args, print_mutation, strip_optimizer, yaml_save)
from utils.loggers import Loggers
# from utils.loggers.wandb.wandb_utils import check_wandb_resume
from utils.loss import ComputeLoss
from utils.metrics import fitness
from utils.plots import plot_evolve, plot_labels
from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
                               smart_resume, torch_distributed_zero_first)

LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
RANK = int(os.getenv('RANK', -1))
# WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
WORLD_SIZE = 1


def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictionary
    save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \
        Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
            opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
    callbacks.run('on_pretrain_routine_start')

    # Directories
    w = save_dir / 'weights'  # weights dir
    (w.parent if evolve else w).mkdir(parents=True, exist_ok=True)  # make dir
    last, best = w / 'last.pt', w / 'best.pt'

    # Hyperparameters
    if isinstance(hyp, str):
        with open(hyp, errors='ignore') as f:
            hyp = yaml.safe_load(f)  # load hyps dict
    LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
    opt.hyp = hyp.copy()  # for saving hyps to checkpoints

    # Save run settings
    if not evolve:
        yaml_save(save_dir / 'hyp.yaml', hyp)
        yaml_save(save_dir / 'opt.yaml', vars(opt))

    # Loggers
    data_dict = None
    if RANK in {-1, 0}:
        loggers = Loggers(save_dir, weights, opt, hyp, LOGGER)  # loggers instance
        if loggers.clearml:
            data_dict = loggers.clearml.data_dict  # None if no ClearML dataset or filled in by ClearML
        if loggers.wandb:
            data_dict = loggers.wandb.data_dict
            if resume:
                weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size

        # Register actions
        for k in methods(loggers):
            callbacks.register_action(k, callback=getattr(loggers, k))

    # Config
    plots = not evolve and not opt.noplots  # create plots
    cuda = device.type != 'cpu'
    init_seeds(opt.seed + 1 + RANK, deterministic=True)
    with torch_distributed_zero_first(LOCAL_RANK):
        data_dict = data_dict or check_dataset(data)  # check if None
    train_path, val_path = data_dict['train'], data_dict['val']
    nc = 1 if single_cls else int(data_dict['nc'])  # number of classes
    names = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names']  # class names
    assert len(names) == nc, f'{len(names)} names found for nc={nc} dataset in {data}'  # check
    is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt')  # COCO dataset

    # Model
    check_suffix(weights, '.pt')  # check weights
    pretrained = weights.endswith('.pt')
    if pretrained:
        with torch_distributed_zero_first(LOCAL_RANK):
            weights = attempt_download(weights)  # download if not found locally
        ckpt = torch.load(weights, map_location='cpu')  # load checkpoint to CPU to avoid CUDA memory leak
        model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
        exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []  # exclude keys
        csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
        csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
        model.load_state_dict(csd, strict=False)  # load
        # print(type(model))
        # exit()
        LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}')  # report
    else:
        model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
    amp = check_amp(model)  # check AMP

    # Freeze
    freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # layers to freeze
    for k, v in model.named_parameters():
        v.requires_grad = True  # train all layers
        # v.register_hook(lambda x: torch.nan_to_num(x))  # NaN to 0 (commented for erratic training results)
        if any(x in k for x in freeze):
            LOGGER.info(f'freezing {k}')
            v.requires_grad = False

    # Image size
    gs = max(int(model.stride.max()), 32)  # grid size (max stride)
    imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)  # verify imgsz is gs-multiple

    # Batch size
    if RANK == -1 and batch_size == -1:  # single-GPU only, estimate best batch size
        batch_size = check_train_batch_size(model, imgsz, amp)
        loggers.on_params_update({"batch_size": batch_size})
    # batch_size = 1

    # Optimizer
    nbs = 64  # nominal batch size
    accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
    hyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decay
    optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])

    # Scheduler
    if opt.cos_lr:
        lf = one_cycle(1, hyp['lrf'], epochs)  # cosine 1->hyp['lrf']
    else:
        lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf']  # linear
    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)  # plot_lr_scheduler(optimizer, scheduler, epochs)

    # EMA
    ema = ModelEMA(model) if RANK in {-1, 0} else None

    # Resume
    best_fitness, start_epoch = 0.0, 0
    if pretrained:
        if resume:
            best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
        del ckpt, csd

    # DP mode
    if cuda and RANK == -1 and torch.cuda.device_count() > 1:
        LOGGER.warning('WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
                       'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')
        model = torch.nn.DataParallel(model)

    # SyncBatchNorm
    if opt.sync_bn and cuda and RANK != -1:
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
        LOGGER.info('Using SyncBatchNorm()')

    # Trainloader

    train_loader, dataset = create_dataloader(train_path,
                                              imgsz,
                                              batch_size // WORLD_SIZE,
                                              gs,
                                              single_cls,
                                              hyp=hyp,
                                              augment=True,
                                              cache=None if opt.cache == 'val' else opt.cache,
                                              rect=opt.rect,
                                              rank=LOCAL_RANK,
                                              workers=workers,
                                              image_weights=opt.image_weights,
                                              quad=opt.quad,
                                              prefix=colorstr('train: '),
                                              shuffle=True)
    labels = np.concatenate(dataset.labels, 0)
    mlc = int(labels[:, 0].max())  # max label class
    assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'

    # Process 0
    if RANK in {-1, 0}:
        val_loader = create_dataloader(val_path,
                                       imgsz,
                                       batch_size // WORLD_SIZE * 2,
                                       gs,
                                       single_cls,
                                       hyp=hyp,
                                       cache=None if noval else opt.cache,
                                       rect=True,
                                       rank=-1,
                                       workers=workers * 2,
                                       pad=0.5,
                                       prefix=colorstr('val: '))[0]
        print(data_dict)
        if not resume:
            if plots:
                plot_labels(labels, names, save_dir)

            # Anchors
            if not opt.noautoanchor:
                check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
            model.half().float()  # pre-reduce anchor precision

        # callbacks.run('on_pretrain_routine_end')

    # DDP mode
    if cuda and RANK != -1:
        model = smart_DDP(model)

    # Model attributes
    nl = de_parallel(model).model[-1].nl  # number of detection layers (to scale hyps)
    hyp['box'] *= 3 / nl  # scale to layers
    hyp['cls'] *= nc / 80 * 3 / nl  # scale to classes and layers
    hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
    hyp['label_smoothing'] = opt.label_smoothing
    model.nc = nc  # attach number of classes to model
    model.hyp = hyp  # attach hyperparameters to model
    model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
    model.names = names

    compute_loss = ComputeLoss(model)  # init loss class

    """
        使用这个脚本来尝试在 Imagenet 数据集上执行量化感知训练
            使用 imagenet 中的数据测试量化精度与 calibration
            默认的 imagenet 数据集位置:Assets/Imagenet_Train, Assets/Imagenet_Valid
            你可以通过软连接创建它们:
                ln -s /home/data/Imagenet/val Assets/Imagenet_Valid
                ln -s /home/data/Imagenet/train Assets/Imagenet_Train
    """

    CFG_DEVICE = 'cuda'  # 一个神奇的字符串，用来确定执行设备
    CFG_BATCHSIZE = 1  # 测试与calib时的 batchsize
    # CFG_INPUT_SHAPE = (CFG_BATCHSIZE, 3, 640, 640)  # 用来确定模型输入的尺寸，好像 imagenet 都是这个尺寸
    # CFG_VALIDATION_DIR = 'Assets/Imagenet_Valid'   # 用来读取 validation dataset
    CFG_VALIDATION_DIR = 'C:/Users/cyc/Desktop/datasets/gundatabase/train'  # 用来读取 validation dataset
    CFG_VALIDATION_DIR = 'C:/Users/cyc/Desktop/datasets/coco128/images/train2017'
    CFG_TRAIN_DIR = 'C:/Users/cyc/Desktop/datasets/gundatabase/train'  # 用来读取 train dataset，注意该集合将被用来 calibrate 你的模型
    # CFG_PLATFORM = TargetPlatform.TRT_INT8  # 用来指定目标平台
    CFG_PLATFORM = TargetPlatform.FPGA_INT8  # 用来指定目标平台

    # ------------------------------------------------------------
    # 在这个例子中我们将向你展示如何在 PPQ 中对你的网络进行量化感知训练
    # 你可以使用带标签的数据执行正常的训练流程，也可以使用类似蒸馏的方式进行无标签训练
    # PPQ 模型的训练过程与 Pytorch 遵循相同的逻辑，你可以使用 Pytorch 中的技巧来获得更好的训练效果
    # ------------------------------------------------------------
    # for name, module in model.named_modules():
    #     # 打印模块的名称和结构
    #     print(f"Module Name: {name}")
    #     print(module)
    #     print("=" * 50)
    # exit()
    graph = load_torch_model(model=model, sample=torch.zeros([1, 3, 640, 640]).cuda())
    # quantizer = FPGAQuantizer(graph=graph)
    # graph = load_onnx_graph('best_1024.onnx')
    quantizer = FPGAQuantizer(graph=graph)
    dispatching = PFL.Dispatcher(graph=graph).dispatch(quantizer.quant_operation_types)

    # PFL.Exporter(platform=TargetPlatform.TRT_INT8).export(
    #     file_path='pattern1_export.onnx', graph=graph, config_path='pattern1_export.json')
    # ------------------------------------------------------------
    # 我们首先进行标准的量化流程，为所有算子初始化量化信息，并进行 Calibration
    # ------------------------------------------------------------
    # for op in graph.operations.values():
    #     print(op.name)
    # exit()
    # Concat_40 往前的所有算子不量化
    from ppq.IR import SearchableGraph
    search_engine = SearchableGraph(graph)
    # for op in search_engine.opset_matching(
    #         sp_expr=lambda x: x.name == 'Concat_122',
    #         rp_expr=lambda x, y: True,
    #         ep_expr=None, direction='up'
    # ):
    #     dispatching[op.name] = TargetPlatform.UNSPECIFIED
    #     print(op.name)
    #
    # # Sigmoid_280 往后的所有算子不量化
    # # Sigmoid_246 往后的所有算子不量化
    # # Sigmoid_314 往后的所有算子不量化
    # for op in search_engine.opset_matching(
    #         sp_expr=lambda x: x.name in {'Sigmoid_516', 'Sigmoid_405', 'Sigmoid_294'},
    #         rp_expr=lambda x, y: True,
    #         ep_expr=None, direction='down'
    # ):
    #     dispatching[op.name] = TargetPlatform.FP32
        # print(op.name)
    # exit()
    # 为算子初始化量化信息
    for op in graph.operations.values():
        quantizer.quantize_operation(
            op_name=op.name, platform=dispatching[op.name])

    # 初始化执行器
    executor = TorchExecutor(graph=graph)
    executor.tracing_operation_meta(inputs=torch.zeros([1, 3, 224, 224]).cuda())
    pipeline = PFL.Pipeline([
        QuantizeSimplifyPass(),
        QuantizeFusionPass(activation_type=quantizer.activation_fusion_types),
        ParameterQuantizePass(),
        RuntimeCalibrationPass(method='kl'),
        PassiveParameterQuantizePass(),
        QuantAlignmentPass(),
        # LearnedStepSizePass(steps=500, block_size=5)
    ])
    pipeline.optimize(
        calib_steps=8, collate_fn=lambda x: x[0].cuda(),
        graph=graph, dataloader=train_loader, executor=executor)

    # PFL.Exporter(platform=TargetPlatform.ONNXRUNTIME).export(
    #     file_path='best_export.onnx', graph=graph, config_path='best_export.json')
    # ------------------------------------------------------------
    # 完成量化后，我们将开始进行 QAT 的模型训练，我们希望你能够注意到：
    #
    # 1. 不能从零开始使用 QAT 的方法完成训练，你应当先训练好浮点的模型，或者在一个预训练的模型基础上进行 QAT finetuning.
    # 2. 你必须完成标准量化流程
    # 3. PPQ Executor 长得很像 Pytorch Module，单机训练应该不会遇到太多困难，但它不支持多卡训练

    # 训练的代码我们封装进了一个叫做 ImageNetTrainer 的东西
    # 你可以打开它看到具体的训练逻辑
    # ------------------------------------------------------------
    trainer = ImageNetTrainer(graph=graph, model=model)
    best_acc = 0

    for epoch in range(10):
        epoch_loss = trainer.epoch(train_loader, compute_loss)
        results = trainer.eval(model.names, val_loader, compute_loss)
        print("epoch_loss:", epoch_loss)
    fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95
    # print(fi)
    # stop = stopper(epoch=epoch, fitness=fi)  # early stop check
    # PFL.Exporter(platform=TargetPlatform.ONNXRUNTIME).export(
    #     file_path='pattern1_export.onnx', graph=graph, config_path='pattern1_export.json')
    if fi > best_fitness:
        best_fitness = fi
        trainer.save("Best.native")
    PFL.Exporter(platform=TargetPlatform.ONNXRUNTIME).export(
        file_path='Output/pattern2_export.onnx', graph=graph, config_path='pattern2_export.json')
    exit(999)

# onnx simplyfied

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
    parser.add_argument('--epochs', type=int, default=300)
    parser.add_argument('--batch-size', type=int, default=1, help='total batch size for all GPUs, -1 for autobatch')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    parser.add_argument('--noval', action='store_true', help='only validate final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
    parser.add_argument('--noplots', action='store_true', help='save no plot files')
    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--workers', type=int, default=0, help='max dataloader workers (per RANK in DDP mode)')
    parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
    parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
    parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
    parser.add_argument('--seed', type=int, default=0, help='Global training seed')
    parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')

    # Weights & Biases arguments
    parser.add_argument('--entity', default=None, help='W&B: Entity')
    parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='W&B: Upload data, "val" option')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
    parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')

    return parser.parse_known_args()[0] if known else parser.parse_args()


def main(opt, callbacks=Callbacks()):
    # Checks
    if RANK in {-1, 0}:
        print_args(vars(opt))
        check_git_status()
        check_requirements()

    # Resume
    if opt.resume and not (check_wandb_resume(opt) or opt.evolve):  # resume from specified or most recent last.pt
        last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
        opt_yaml = last.parent.parent / 'opt.yaml'  # train options yaml
        opt_data = opt.data  # original dataset
        if opt_yaml.is_file():
            with open(opt_yaml, errors='ignore') as f:
                d = yaml.safe_load(f)
        else:
            d = torch.load(last, map_location='cpu')['opt']
        opt = argparse.Namespace(**d)  # replace
        opt.cfg, opt.weights, opt.resume = '', str(last), True  # reinstate
        if is_url(opt_data):
            opt.data = check_file(opt_data)  # avoid HUB resume auth timeout
    else:
        opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
            check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project)  # checks
        assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
        if opt.evolve:
            if opt.project == str(ROOT / 'runs/train'):  # if default project name, rename to runs/evolve
                opt.project = str(ROOT / 'runs/evolve')
            opt.exist_ok, opt.resume = opt.resume, False  # pass resume to exist_ok and disable resume
        if opt.name == 'cfg':
            opt.name = Path(opt.cfg).stem  # use model.yaml as name
        opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))

    # DDP mode
    device = select_device(opt.device, batch_size=opt.batch_size)
    if LOCAL_RANK != -1:
        msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
        assert not opt.image_weights, f'--image-weights {msg}'
        assert not opt.evolve, f'--evolve {msg}'
        assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
        assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
        assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
        torch.cuda.set_device(LOCAL_RANK)
        device = torch.device('cuda', LOCAL_RANK)
        dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")

    # Train
    if not opt.evolve:
        tr, va = train(opt.hyp, opt, device, callbacks)

    return tr, va
    # Evolve hyperparameters (optional)
    # else:
    #     # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
    #     meta = {
    #         'lr0': (1, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
    #         'lrf': (1, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
    #         'momentum': (0.3, 0.6, 0.98),  # SGD momentum/Adam beta1
    #         'weight_decay': (1, 0.0, 0.001),  # optimizer weight decay
    #         'warmup_epochs': (1, 0.0, 5.0),  # warmup epochs (fractions ok)
    #         'warmup_momentum': (1, 0.0, 0.95),  # warmup initial momentum
    #         'warmup_bias_lr': (1, 0.0, 0.2),  # warmup initial bias lr
    #         'box': (1, 0.02, 0.2),  # box loss gain
    #         'cls': (1, 0.2, 4.0),  # cls loss gain
    #         'cls_pw': (1, 0.5, 2.0),  # cls BCELoss positive_weight
    #         'obj': (1, 0.2, 4.0),  # obj loss gain (scale with pixels)
    #         'obj_pw': (1, 0.5, 2.0),  # obj BCELoss positive_weight
    #         'iou_t': (0, 0.1, 0.7),  # IoU training threshold
    #         'anchor_t': (1, 2.0, 8.0),  # anchor-multiple threshold
    #         'anchors': (2, 2.0, 10.0),  # anchors per output grid (0 to ignore)
    #         'fl_gamma': (0, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
    #         'hsv_h': (1, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
    #         'hsv_s': (1, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
    #         'hsv_v': (1, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
    #         'degrees': (1, 0.0, 45.0),  # image rotation (+/- deg)
    #         'translate': (1, 0.0, 0.9),  # image translation (+/- fraction)
    #         'scale': (1, 0.0, 0.9),  # image scale (+/- gain)
    #         'shear': (1, 0.0, 10.0),  # image shear (+/- deg)
    #         'perspective': (0, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
    #         'flipud': (1, 0.0, 1.0),  # image flip up-down (probability)
    #         'fliplr': (0, 0.0, 1.0),  # image flip left-right (probability)
    #         'mosaic': (1, 0.0, 1.0),  # image mixup (probability)
    #         'mixup': (1, 0.0, 1.0),  # image mixup (probability)
    #         'copy_paste': (1, 0.0, 1.0)}  # segment copy-paste (probability)
    #
    #     with open(opt.hyp, errors='ignore') as f:
    #         hyp = yaml.safe_load(f)  # load hyps dict
    #         if 'anchors' not in hyp:  # anchors commented in hyp.yaml
    #             hyp['anchors'] = 3
    #     if opt.noautoanchor:
    #         del hyp['anchors'], meta['anchors']
    #     opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir)  # only val/save final epoch
    #     # ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indices
    #     evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
    #     if opt.bucket:
    #         os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {evolve_csv}')  # download evolve.csv if exists
    #
    #     for _ in range(opt.evolve):  # generations to evolve
    #         if evolve_csv.exists():  # if evolve.csv exists: select best hyps and mutate
    #             # Select parent(s)
    #             parent = 'single'  # parent selection method: 'single' or 'weighted'
    #             x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
    #             n = min(5, len(x))  # number of previous results to consider
    #             x = x[np.argsort(-fitness(x))][:n]  # top n mutations
    #             w = fitness(x) - fitness(x).min() + 1E-6  # weights (sum > 0)
    #             if parent == 'single' or len(x) == 1:
    #                 # x = x[random.randint(0, n - 1)]  # random selection
    #                 x = x[random.choices(range(n), weights=w)[0]]  # weighted selection
    #             elif parent == 'weighted':
    #                 x = (x * w.reshape(n, 1)).sum(0) / w.sum()  # weighted combination
    #
    #             # Mutate
    #             mp, s = 0.8, 0.2  # mutation probability, sigma
    #             npr = np.random
    #             npr.seed(int(time.time()))
    #             g = np.array([meta[k][0] for k in hyp.keys()])  # gains 0-1
    #             ng = len(meta)
    #             v = np.ones(ng)
    #             while all(v == 1):  # mutate until a change occurs (prevent duplicates)
    #                 v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
    #             for i, k in enumerate(hyp.keys()):  # plt.hist(v.ravel(), 300)
    #                 hyp[k] = float(x[i + 7] * v[i])  # mutate
    #
    #         # Constrain to limits
    #         for k, v in meta.items():
    #             hyp[k] = max(hyp[k], v[1])  # lower limit
    #             hyp[k] = min(hyp[k], v[2])  # upper limit
    #             hyp[k] = round(hyp[k], 5)  # significant digits
    #
    #         # Train mutation
    #         results = train(hyp.copy(), opt, device, callbacks)
    #         callbacks = Callbacks()
    #         # Write mutation results
    #         print_mutation(results, hyp.copy(), save_dir, opt.bucket)
    #
    #     # Plot results
    #     plot_evolve(evolve_csv)
    #     LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'
    #                 f"Results saved to {colorstr('bold', save_dir)}\n"
    #                 f'Usage example: $ python train.py --hyp {evolve_yaml}')


def run(**kwargs):
    # Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
    opt = parse_opt(True)
    for k, v in kwargs.items():
        setattr(opt, k, v)
    main(opt)
    return opt


if __name__ == "__main__":
    # torch.use_deterministic_algorithms(True)
    opt = parse_opt()
    main(opt)
